import numpy as np
import matplotlib.pyplot as plt
import rasterio
from sklearn.utils import resample
from scipy.stats import linregress, binned_statistic
import os


def read_geotiff_with_nodata(file_path):
    """读取GeoTIFF文件，处理NoData值"""
    with rasterio.open(file_path) as src:
        data = src.read(1)  # 读取第一个波段
        nodata = src.nodata  # 获取NoData值
        mask = (data == nodata) if nodata is not None else np.zeros_like(data, dtype=bool)
        return data, mask, src.tags().get('LAYER_NAME', os.path.basename(file_path))


def prepare_data(file1, file2, sample_ratio=1.0, random_state=None, min_file2=0.1):
    """
    准备绘图数据，返回采样后的数据对

    参数:
        file1, file2: GeoTIFF文件路径
        sample_ratio: 采样比例(0-1)
        random_state: 随机种子
        min_file2: file2的最小有效值
    返回:
        (sampled_data1, sampled_data2), (layer1, layer2)
    """
    # 读取数据
    data1, mask1, layer1 = read_geotiff_with_nodata(file1)
    data2, mask2, layer2 = read_geotiff_with_nodata(file2)

    # 校验尺寸
    assert data1.shape == data2.shape, "图像尺寸不一致"

    # 合并掩码（NoData或file2<min_file2）
    combined_mask = mask1 | mask2 | (data2 < min_file2)
    valid_coords = np.column_stack(np.where(~combined_mask))

    # 采样
    n_samples = int(len(valid_coords) * sample_ratio)
    if n_samples == 0:
        raise ValueError("无有效数据点")

    if sample_ratio < 1.0:
        sampled_coords = resample(valid_coords, replace=False,
                                  n_samples=n_samples, random_state=random_state)
    else:
        sampled_coords = valid_coords

    return (data1[sampled_coords[:, 0], sampled_coords[:, 1]],
            data2[sampled_coords[:, 0], sampled_coords[:, 1]]), (layer1, layer2)


def plot_scatter(data_pair, layers):
    """绘制散点图"""
    data1, data2 = data_pair
    layer1, layer2 = layers

    plt.figure(figsize=(10, 8))
    plt.scatter(data1, data2, alpha=0.5, s=5)

    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimSun']
    plt.rcParams['axes.unicode_minus'] = False

    plt.title(f'{layer1} vs {layer2} 散点图', fontsize=14)
    plt.xlabel(layer1, fontsize=12)
    plt.ylabel(layer2, fontsize=12)
    plt.grid(True, linestyle='--', alpha=0.6)
    plt.tight_layout()
    plt.show()


def plot_boxplot(data_pair, layers, n_bins=100, label_interval=10, output_csv=None):
    """
    绘制分段箱线图并保存统计结果为CSV

    参数:
        data_pair: (data1, data2) 数据对
        layers: (layer1_name, layer2_name) 图层名称
        n_bins: 分段数量
        label_interval: x轴标签显示间隔
        output_csv: CSV输出路径 (None则不保存)
    """
    import pandas as pd
    data1, data2 = data_pair
    layer1, layer2 = layers

    # 计算分段统计
    bin_means, bin_edges, _ = binned_statistic(data1, data2,
                                               statistic='mean',
                                               bins=n_bins)
    bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2

    # 准备箱线图数据
    box_data = []
    valid_bin_centers = []
    valid_bin_means = []
    for i in range(len(bin_edges) - 1):
        mask = (data1 >= bin_edges[i]) & (data1 < bin_edges[i + 1])
        if np.any(mask):
            box_data.append(data2[mask])
            valid_bin_centers.append(bin_centers[i])
            valid_bin_means.append(bin_means[i])

    # 保存为CSV
    if output_csv:
        df = pd.DataFrame({
            f'{layer1}_bin_center': valid_bin_centers,
            f'{layer1}_bin_range': [f"{bin_edges[i]:.2f}-{bin_edges[i + 1]:.2f}"
                                    for i in range(len(bin_edges) - 1) if
                                    np.any((data1 >= bin_edges[i]) & (data1 < bin_edges[i + 1]))],
            f'{layer2}_mean': valid_bin_means,
            'count': [len(b) for b in box_data]
        })
        df.to_csv(output_csv, index=False, encoding='utf_8_sig')  # 支持中文

    # 绘图
    plt.figure(figsize=(12, 6))

    # 箱线图样式设置
    flierprops = dict(marker='o', markersize=0, linestyle='none')   # 设置异常点的样式及大小1像素
    boxplot = plt.boxplot(box_data, positions=valid_bin_centers,
                          widths=(bin_edges[1] - bin_edges[0]) * 0.8,
                          showmeans=True, patch_artist=True,
                          flierprops=flierprops)

    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimSun']
    plt.rcParams['axes.unicode_minus'] = False

    # 智能x轴标签显示
    if len(valid_bin_centers) > label_interval:
        xticks_pos = valid_bin_centers[::label_interval]
        xticks_labels = [f"{bin_edges[i]:.1f}" for i in range(0, len(bin_edges) - 1, label_interval)]
    else:
        xticks_pos = valid_bin_centers
        xticks_labels = [f"{bin_edges[i]:.1f}" for i in range(len(bin_edges) - 1)]

    plt.xticks(xticks_pos, xticks_labels, rotation=45, ha='right')

    plt.title(f'{layer1} 分段箱线图 ({n_bins}段)', fontsize=14)
    plt.xlabel(f'{layer1} 值域范围 (每{label_interval}段显示)', fontsize=12)
    plt.ylabel(f'{layer2} 数值分布', fontsize=12)
    plt.grid(True, linestyle='--', alpha=0.6)
    plt.tight_layout()
    plt.show()


# 示例用法
if __name__ == "__main__":
    file1 = r"D:\GPPvsWater\work\散点图\AI.tif"
    file2 = r"D:\GPPvsWater\work\散点图\ndvimean_con.tif"

    try:
        # 准备数据
        data_pair, layers = prepare_data(file1, file2,
                                         sample_ratio=0.3,
                                         random_state=42,
                                         min_file2=0.1)

        # 分别绘图
        # plot_scatter(data_pair, layers)
        plot_boxplot(data_pair, layers, n_bins=100, label_interval=10, output_csv="./boxplot_stats.csv")

    except Exception as e:
        print(f"发生错误: {e}")